
#include "nll_loss_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"

constexpr int32_t CORENUM = 20;

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context) {
  NLLLossTilingData tiling;
  const gert::StorageShape* x1_shape = context->GetInputShape(0);
  const gert::StorageShape* xtarget_shape = context->GetInputShape(0);

  const char* reductionStr = context->GetAttrs()->GetStr(0);
  const int64_t* ignore_index_addr = context->GetAttrs()->GetInt(1);
  uint32_t reduction = -1;
  uint32_t ignore_index = *ignore_index_addr;
  if (strcmp(reductionStr, "sum") == 0) {
    reduction = 1;
  } else if (strcmp(reductionStr, "mean") == 0) {
    reduction = 2;
  }

  uint32_t Dim = x1_shape->GetStorageShape().GetDimNum();

  uint32_t sizeN, sizeC;

  if (Dim == 2) {
    sizeN = x1_shape->GetStorageShape().GetDim(0);
    sizeC = x1_shape->GetStorageShape().GetDim(1);
  } else {
    sizeN = 1;
    sizeC = x1_shape->GetStorageShape().GetDim(0);
  }


  if (Dim == 1) {
    context->SetTilingKey(1);
  } else if (Dim == 2){
    if (reduction==2)
    {
      context->SetTilingKey(2);
    } else
    {
      context->SetTilingKey(3);
    }
  } else {
    context->SetTilingKey(4);
  }

  tiling.set_dimNum(Dim);
  tiling.set_reduction(reduction);
  tiling.set_ignore_index(ignore_index);
  tiling.set_sizeN(sizeN);  // input_X shape = {N,C}, where C need to reduce
  tiling.set_sizeC(sizeC);

  // -------- set cores tiling --------
  uint32_t coreNum = (sizeN < CORENUM) ? sizeN : CORENUM;
  // printf("coreNum %d %d\n", coreNum, sizeN);

  uint32_t coreRowNum  = sizeN / coreNum;
  uint32_t coreRowMod = sizeN % coreNum;

  tiling.set_coreNum(coreNum);
  tiling.set_coreRowNum(coreRowNum);
  tiling.set_coreRowMod(coreRowMod);

  context->SetBlockDim(coreNum);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                      context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  return ge::GRAPH_SUCCESS;

  // size_t usrSize = (256) * coreNum;  // 设置用户需要使用的workspace大小。
  // // 如需要使用系统workspace需要调用GetLibApiWorkSpaceSize获取系统workspace的大小。
  // auto ascendcPlatform =
  //     platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  // uint32_t sysWorkspaceSize = ascendcPlatform.GetLibApiWorkSpaceSize();
  // size_t* currentWorkspace = context->GetWorkspaceSizes(1);
  // currentWorkspace[0] = usrSize + sysWorkspaceSize;
  // return ge::GRAPH_SUCCESS;
}
}  // namespace optiling

// namespace ge {
// static ge::graphStatus InferShape(gert::InferShapeContext* context)
// {
//     const gert::Shape* x1_shape = context->GetInputShape(0);
//     gert::Shape* y_shape = context->GetOutputShape(0);
//     *y_shape = *x1_shape;
//     return GRAPH_SUCCESS;
// }
// }

namespace ops {
class NLLLoss : public OpDef {
 public:
  explicit NLLLoss(const char* name) : OpDef(name) {
    this->Input("x")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
    this->Input("target")
        .ParamType(REQUIRED)
        .DataType({ge::DT_INT32})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
    this->Input("weight")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
    this->Output("y")
        .ParamType(REQUIRED)
        .DataType({ge::DT_FLOAT})
        .Format({ge::FORMAT_ND})
        .UnknownShapeFormat({ge::FORMAT_ND});
    this->Attr("reduction").AttrType(OPTIONAL).String("mean");
    this->Attr("ignore_index").AttrType(OPTIONAL).Int(-100);
    // this->SetInferShape(ge::InferShape);

    this->AICore().SetTiling(optiling::TilingFunc);
    this->AICore().AddConfig("ascend910b");
    // this->AICore().AddConfig("ascend310b");
  }
};

OP_ADD(NLLLoss);
}  // namespace ops
